package clrs

import (
	"fmt"
	"strings"
)

/**
 * 红黑树，
 * 性质：
 * 1. 每个结点或是红的，或是黑的。
 * 2. 根结点是黑的NewRedBlackTreeNode
 * 3. 每个空(NilNode)结点是黑的
 * 4. 如果一个结点是红的，则它的两个儿子都是黑的。
 * 5. 对每个结点，从该结点到其子孙结点的所有路径上包含相同数目的黑结点。
 */
type RedBlackTree struct {
	Root *RedBlackTreeNode
}

var NilNode = NewRedBlackTreeNode(0, 0, nil)

func NewRedBlackTree() *RedBlackTree {
	that := RedBlackTree{
		Root: NilNode,
	}

	return &that
}

/**
 * 左旋转，将一个结点置换到它的左子结点的位置。
 *
 * 需要维护二叉查找树的性质（左小右大）。
 *
 * leftRotate(x) 和 rightRotate(right)互为逆操作。
 *
 * ```txt
 * left < x < a < right < b
 *
 * From
 * x
 * |-> left
 * |-> right
 * |   |-> a
 * |   |-> b
 *
 * To
 * right
 * |-> x
 * |   |-> left
 * |   |-> a
 * |-> b
 * ```
 *
 * @param node the node
 */
func (that *RedBlackTree) leftRotate(node *RedBlackTreeNode) {
	if node == NilNode {
		return
	}
	originalRight := node.RightChild

	// 结点的右子结点 <- 原来右节点的左子结点。
	node.RightChild = originalRight.LeftChild

	// 原来的右节点的左子节点 <- 原结点。
	if originalRight.LeftChild != NilNode {
		originalRight.LeftChild.Parent = node
	}

	// 原来的右节点置换到原结点的位置。
	originalRight.Parent = node.Parent

	// 调整原父节点。
	if node.Parent == NilNode {
		// case 原结点是根节点，
		that.Root = originalRight
	} else if node == node.Parent.LeftChild {
		// case 原结点是父节点的左子节点。
		node.Parent.LeftChild = originalRight
	} else {
		// case 原结点是父节点的右子节点。
		node.Parent.RightChild = originalRight
	}

	originalRight.LeftChild = node
	node.Parent = originalRight
}

/**
 * 右旋转，将一个结点置换到它的右子结点的位置。
 *
 * 需要维护二叉查找树的性质（左小右大）。
 *
 * rightRotate(x) 和 leftRotate(left)互为逆操作。
 *
 * ```txt
 * a < left < b < x < right
 *
 * From
 * x
 * |-> left
 * |   |-> a
 * |   |-> b
 * |-> right
 *
 * To
 * left
 * |-> a
 * |-> x
 * |   |-> b
 * |   |-> right
 * ```
 * @param node the node
 */
func (that *RedBlackTree) rightRotate(node *RedBlackTreeNode) {
	if node == NilNode {
		return
	}
	var originalLeft = node.LeftChild

	// 结点的左子结点 <- 原来左节点的右子结点。
	node.LeftChild = originalLeft.RightChild

	// 原来的左节点的右子节点 <- 原结点。
	if originalLeft.RightChild != NilNode {
		originalLeft.RightChild.Parent = node
	}

	// 原来的左节点置换到原结点的位置。
	originalLeft.Parent = node.Parent

	// 调整原父节点。
	if node.Parent == NilNode {
		// case 原结点是根节点，
		that.Root = originalLeft
	} else if isLeftChild(node) {
		// case 原结点是父节点的左子节点。
		node.Parent.LeftChild = originalLeft
	} else {
		// case 原结点是父节点的右子节点。
		node.Parent.RightChild = originalLeft
	}

	originalLeft.RightChild = node
	node.Parent = originalLeft
}

func (that *RedBlackTree) Delete(node *RedBlackTreeNode) *RedBlackTreeNode {
	if node == NilNode {
		return NilNode
	}

	var y *RedBlackTreeNode = NilNode
	var x *RedBlackTreeNode = NilNode
	if node.LeftChild == NilNode || node.RightChild == NilNode {
		y = node
	} else {
		y = that.successor(node)
	}

	if y.LeftChild != NilNode {
		x = y.LeftChild
	} else {
		x = y.RightChild
	}

	if x != NilNode {
		x.Parent = y.Parent
	}

	if y.Parent == NilNode {
		that.Root = x
	} else if isLeftChild(y) {
		y.Parent.LeftChild = x
	} else {
		y.Parent.RightChild = x
	}

	if y != node {
		y.CopyAttributesTo(node)
	}

	if y.IsBlack() {
		that.deleteFixUp(x)
	}

	return y
}

func (that *RedBlackTree) deleteFixUp(node *RedBlackTreeNode) {
	if node == NilNode {
		return
	}
	for node != that.Root && node.IsBlack() {
		if isLeftChild(node) {
			var w = node.Parent.RightChild
			if w.isRed {
				w.SetBlack(true)
				node.Parent.SetRed(true)
				that.leftRotate(node.Parent)
				w = node.Parent.RightChild
			}
			if w.LeftChild.IsBlack() && w.RightChild.IsBlack() {
				w.SetRed(true)
				node = node.Parent
			} else {
				if w.RightChild.IsBlack() {
					w.LeftChild.SetBlack(true)
					w.SetRed(true)
					that.rightRotate(w)
					w = node.Parent.RightChild
				}
				w.isRed = node.Parent.isRed
				node.Parent.SetBlack(true)
				w.RightChild.SetBlack(true)
				that.leftRotate(node.Parent)
				node = that.Root
			}
		} else {
			// same as then clause with "right" and "left" exchanged
			var w = node.Parent.LeftChild
			if w.isRed {
				w.SetBlack(true)
				node.Parent.SetRed(true)
				that.rightRotate(node.Parent)
				w = node.Parent.LeftChild
			}
			if w.RightChild.IsBlack() && w.LeftChild.IsBlack() {
				w.SetRed(true)
				node = node.Parent
			} else {
				if w.LeftChild.IsBlack() {
					w.RightChild.SetBlack(true)
					w.SetRed(true)
					that.leftRotate(w)
					w = node.Parent.LeftChild
				}
				w.isRed = node.Parent.isRed
				node.Parent.SetBlack(true)
				w.LeftChild.SetBlack(true)
				that.rightRotate(node.Parent)
				node = that.Root
			}
		}
	}
	node.SetBlack(true)
}

func (that *RedBlackTree) Insert(key int, value int) {
	// create the node and configure parent properties
	var node = NewRedBlackTreeNode(key, value, NilNode)

	if that.Root == NilNode {
		that.Root = node
	} else {
		// Find the insert parent
		var targetParent *RedBlackTreeNode = NilNode
		var x = that.Root
		for x != NilNode {
			targetParent = x
			if key < x.Key {
				x = x.LeftChild
			} else {
				x = x.RightChild
			}
		}

		node.Parent = targetParent
		if node.Key < targetParent.Key {
			targetParent.LeftChild = node
		} else {
			targetParent.RightChild = node
		}
	}

	// red
	node.SetRed(true)
	that.insertFixUp(node)
}

/**
 * 保证红黑树的性质
 * @param node the node
 */
func (that *RedBlackTree) insertFixUp(node *RedBlackTreeNode) {
	// 由于node是红色，所以父是红色违反了规则
	for node.Parent.isRed {
		if isLeftChild(node.Parent) {
			var y = node.Parent.Parent.RightChild
			if y.isRed {
				// case 1
				node.Parent.SetBlack(true)
				y.SetBlack(true)
				if node.Parent.Parent != NilNode {
					node.Parent.Parent.SetRed(true)
				}
				node = node.Parent.Parent
			} else {
				// case 2 or 3
				if isRightChild(node) {
					// case 2 node 为右子
					node = node.Parent
					that.leftRotate(node)
				}
				// case 3 node 为左子
				node.Parent.SetBlack(true)
				node.Parent.Parent.SetRed(true)
				that.rightRotate(node.Parent.Parent)
			}
		} else {
			var y = node.Parent.Parent.LeftChild
			if y.isRed {
				node.Parent.SetBlack(true)
				y.SetBlack(true)
				if node.Parent.Parent != NilNode {
					node.Parent.Parent.SetRed(true)
				}
				node = node.Parent.Parent
			} else {
				if isLeftChild(node) {
					node = node.Parent
					that.rightRotate(node)
				}
				node.Parent.SetBlack(true)
				node.Parent.Parent.SetRed(true)
				that.leftRotate(node.Parent.Parent)
			}
		}
	}
	that.Root.SetBlack(true)
}

func isLeftChild(node *RedBlackTreeNode) bool {
	return node == node.Parent.LeftChild
}

func isRightChild(node *RedBlackTreeNode) bool {
	return node == node.Parent.RightChild
}

func (that *RedBlackTree) Tree(node *RedBlackTreeNode, level int) string {
	if node == nil {
		node = that.Root
	}

	var result string = ""
	// output level
	var levelText string = ""
	if level > 0 {
		if level > 1 {
			levelText = strings.Repeat("│   ", level-1)
		}
		levelText += "├── "
	}

	// output node
	result = levelText + that.renderTreeNodeColor(node)

	// case null node -> do not output children
	if node == NilNode {
		return result
	}

	// output children
	if node.LeftChild != NilNode || node.RightChild != NilNode {
		result += "\n" + that.Tree(node.LeftChild, level+1)
		result += "\n" + that.Tree(node.RightChild, level+1)
	}

	return result
}

func (that *RedBlackTree) renderTreeNodeColor(node *RedBlackTreeNode) string {
	if node == NilNode {
		return "(nil)"
	}

	if node.IsBlack() {
		return fmt.Sprint(node.value)
	}

	return "\x1b[31m" + fmt.Sprint(node.value) + "\x1b[0m"
}

func (that *RedBlackTree) Maximum(node *RedBlackTreeNode) *RedBlackTreeNode {
	if node == NilNode {
		return node
	}

	for node.RightChild != NilNode {
		node = node.RightChild
	}

	return node
}

func (that *RedBlackTree) Minimum(node *RedBlackTreeNode) *RedBlackTreeNode {
	if node == NilNode {
		return node
	}

	for node.LeftChild != NilNode {
		node = node.LeftChild
	}

	return node
}

/**
 * Search 的非递归版本
 * @param key the key
 */
func (that *RedBlackTree) Search(key int) *RedBlackTreeNode {
	var node = that.Root
	for node != NilNode && node.Key != key {
		if key < node.Key {
			node = node.LeftChild
		} else {
			node = node.RightChild
		}
	}

	if key == node.Key {
		return node
	}
	return nil
}

/**
 * 返回指定结点在中序遍历中的后继结点
 * @param node the node
 */
func (that *RedBlackTree) successor(node *RedBlackTreeNode) *RedBlackTreeNode {
	if node.RightChild != NilNode {
		return that.Minimum(node.RightChild)
	}

	var y = node.Parent
	for y != NilNode && node == y.RightChild {
		node = y
		y = y.Parent
	}

	return y
}

/**
 * 返回指定索引的结点。
 * 也是第最小i个的结点。
 * @param index the index, return index-nd node
 * @param node the node, optional the root node
 */
func (that *RedBlackTree) Get(index int32, node *RedBlackTreeNode) *RedBlackTreeNode {
	var r = node.LeftChild.size + 1
	if index == r {
		return node
	}

	if index < r {
		return that.Get(index, node.LeftChild)
	}

	return that.Get(index, node.RightChild)
}

/**
 * 返回指定结点的顺序排名。
 * @param node the node, optional the root node
 */
func (that *RedBlackTree) GetRank(node *RedBlackTreeNode) int32 {
	var rank = node.LeftChild.size + 1
	var tnode *RedBlackTreeNode = node
	for tnode == that.Root {
		if isRightChild(tnode) {
			rank += tnode.Parent.LeftChild.size + 1
		}
		tnode = tnode.Parent
	}

	return rank
}

/**
 * 返回指定结点在中序遍历中的前序结点
 * @param node the node
 */
func (that *RedBlackTree) Predecessor(node *RedBlackTreeNode) *RedBlackTreeNode {
	if node.LeftChild != NilNode {
		return that.Maximum(node.LeftChild)
	}

	var y = node.Parent
	for y != NilNode && node == y.LeftChild {
		node = y
		y = y.Parent
	}

	return y
}

func (that *RedBlackTree) WalkThrough(
	node *RedBlackTreeNode,
	callback func(node *RedBlackTreeNode)) {
	if node.LeftChild != NilNode {
		that.WalkThrough(node.LeftChild, callback)
	}
	callback(node)
	if node.RightChild != NilNode {
		that.WalkThrough(node.RightChild, callback)
	}
}

func (that *RedBlackTree) Validate() []string {
	var errors = []string{}
	var err = that.ValidateRootIsBlack()
	if len(err) > 0 {
		errors = append(errors, err...)
	}

	err = that.ValidateLeftSmallerRightBigger()
	if len(err) > 0 {
		errors = append(errors, err...)
	}

	err = that.ValidateRedNodeHasTwoBlack()
	if len(err) > 0 {
		errors = append(errors, err...)
	}

	// return
	return errors
}

func (that *RedBlackTree) ValidateLeftSmallerRightBigger() []string {
	var errors = []string{}
	that.WalkThrough(that.Root, func(node *RedBlackTreeNode) {
		if node == NilNode {
			return
		}
		if node.LeftChild != NilNode {
			if node.LeftChild.Key > node.Key {
				errors = append(errors, `the left node of (${node.Key}) must be small`)
			}
		}
		if node.RightChild != NilNode {
			if node.RightChild.Key < node.Key {
				errors = append(errors, `the right node of (${node.Key}) must be bigger`)
			}
		}
	})
	return errors
}

func (that *RedBlackTree) ValidateRootIsBlack() []string {
	if that.Root.isRed {
		return []string{"The root node must be black"}
	}
	return []string{}
}

func (that *RedBlackTree) ValidateRedNodeHasTwoBlack() []string {
	var errors = []string{}
	that.WalkThrough(that.Root, func(node *RedBlackTreeNode) {
		if node == NilNode || node.IsBlack() {
			return
		}
		if node.LeftChild.isRed {
			errors = append(errors, `the left node of (${node.Key}) must be black`)
		}
		if node.RightChild.isRed {
			errors = append(errors, `the right node of (${node.Key}) must be black`)
		}
	})
	return errors
}
